import os
from flask import Flask
from extensions import db
from datetime import datetime
from sqlalchemy import func

def init_database(app):
    """初始化数据库"""
    # 数据库配置
    basedir = os.path.abspath(os.path.dirname(__file__))
    
    # 使用SQLite数据库（生产环境可以改为PostgreSQL）
    app.config['SQLALCHEMY_DATABASE_URI'] = os.environ.get('DATABASE_URL') or \
        'sqlite:///' + os.path.join(basedir, 'biomni.db')
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
    app.config['SQLALCHEMY_ENGINE_OPTIONS'] = {
        'pool_pre_ping': True,
        'pool_recycle': 300,
    }
    
    # 初始化数据库
    db.init_app(app)
    
    with app.app_context():
        # 创建所有表
        db.create_all()
        
        # 初始化默认数据
        init_default_data()

def init_default_data():
    """初始化默认数据"""
    # 导入模型类
    from models import User, ModelPricing, UserRole, UserStats, ConversationFile
    
    # 创建默认管理员用户
    admin_username = os.environ.get('ADMIN_USERNAME', 'admin')
    admin_user = User.query.filter_by(username=admin_username).first()
    if not admin_user:
        admin_user = User(
            username=admin_username,
            email=os.environ.get('ADMIN_EMAIL', 'admin@biomni.com'),
            role=UserRole.ADMIN,
            balance=float(os.environ.get('ADMIN_BALANCE', '1000.0')),
            weekly_request_limit=int(os.environ.get('ADMIN_WEEKLY_LIMIT', '10000')),
            is_active=True,
            is_verified=True
        )
        admin_password = os.environ.get('ADMIN_PASSWORD', 'admin123')
        admin_user.set_password(admin_password)
        db.session.add(admin_user)
    
    # 初始化模型定价
    pricing_data = [
        {
            'model_name': 'gpt-4o',
            'platform': 'openai',
            'input_price_per_1k': 0.005,   # $0.005 per 1k input tokens
            'output_price_per_1k': 0.015   # $0.015 per 1k output tokens
        },
        {
            'model_name': 'gpt-4o-mini',
            'platform': 'openai',
            'input_price_per_1k': 0.00015,  # $0.00015 per 1k input tokens
            'output_price_per_1k': 0.0006   # $0.0006 per 1k output tokens
        },
        {
            'model_name': 'gpt-4.1',
            'platform': 'openai',
            'input_price_per_1k': 0.03,     # $0.03 per 1k input tokens
            'output_price_per_1k': 0.06     # $0.06 per 1k output tokens
        },
        {
            'model_name': 'claude-sonnet-4-20250514',
            'platform': 'anthropic',
            'input_price_per_1k': 0.003,    # $0.003 per 1k input tokens
            'output_price_per_1k': 0.015    # $0.015 per 1k output tokens
        },
        {
            'model_name': 'claude-3-5-sonnet-20241022',
            'platform': 'anthropic',
            'input_price_per_1k': 0.003,    # $0.003 per 1k input tokens
            'output_price_per_1k': 0.015    # $0.015 per 1k output tokens
        },
        {
            'model_name': 'claude-3-5-haiku-20241022',
            'platform': 'anthropic',
            'input_price_per_1k': 0.001,    # $0.001 per 1k input tokens
            'output_price_per_1k': 0.005    # $0.005 per 1k output tokens
        },
        {
            'model_name': 'gpt-3.5-turbo',
            'platform': 'openai',
            'input_price_per_1k': 0.001,    # $0.001 per 1k input tokens
            'output_price_per_1k': 0.002    # $0.002 per 1k output tokens
        },
        {
            'model_name': 'simple_response',
            'platform': 'internal',
            'input_price_per_1k': 0.0001,   # $0.0001 per 1k input tokens (低成本简单对话)
            'output_price_per_1k': 0.0002   # $0.0002 per 1k output tokens (低成本简单对话)
        }
    ]
    
    for pricing_info in pricing_data:
        existing_pricing = ModelPricing.query.filter_by(
            model_name=pricing_info['model_name']
        ).first()
        
        if not existing_pricing:
            pricing = ModelPricing(**pricing_info)
            db.session.add(pricing)
    
    # 提交所有更改
    try:
        db.session.commit()
        print("默认数据初始化完成")
    except Exception as e:
        db.session.rollback()
        print(f"默认数据初始化失败: {e}")

def create_test_user():
    """创建测试用户（开发环境使用）"""
    from models import User, UserRole
    
    test_user = User.query.filter_by(username='testuser').first()
    if not test_user:
        test_user = User(
            username='testuser',
            email='test@example.com',
            role=UserRole.USER,
            balance=100.0,  # 给测试用户100美元余额
            weekly_request_limit=100,
            is_active=True,
            is_verified=True
        )
        test_user.set_password('test123')
        db.session.add(test_user)
        
        try:
            db.session.commit()
            print("测试用户创建成功")
        except Exception as e:
            db.session.rollback()
            print(f"测试用户创建失败: {e}")

def get_user_stats(user_id):
    """获取用户统计信息"""
    from models import User, RequestRecord, CostRecord
    from sqlalchemy import func
    
    user = User.query.get(user_id)
    if not user:
        return None
    
    # 统计信息
    from models import RequestRecord
    total_requests = db.session.query(RequestRecord).filter_by(user_id=user_id).count()
    completed_requests = db.session.query(RequestRecord).filter_by(
        user_id=user_id, status='completed'
    ).count()
    
    total_tokens = db.session.query(
        func.sum(RequestRecord.total_tokens)
    ).filter_by(user_id=user_id).scalar() or 0
    
    total_cost = db.session.query(
        func.sum(CostRecord.cost)
    ).filter_by(user_id=user_id).scalar() or 0.0
    
    # 本周统计
    weekly_requests = user.get_weekly_request_count()
    
    return {
        'user': user.to_dict(),
        'stats': {
            'total_requests': total_requests,
            'completed_requests': completed_requests,
            'weekly_requests': weekly_requests,
            'total_tokens': total_tokens,
            'total_cost': round(total_cost, 4),
            'remaining_weekly_requests': max(0, user.weekly_request_limit - weekly_requests)
        }
    }

def update_model_pricing(model_name, input_price, output_price):
    """更新模型定价"""
    from models import ModelPricing
    
    pricing = ModelPricing.query.filter_by(model_name=model_name).first()
    if pricing:
        pricing.input_price_per_1k = input_price
        pricing.output_price_per_1k = output_price
        pricing.updated_at = datetime.utcnow()
    else:
        # 如果不存在，创建新的定价记录
        pricing = ModelPricing(
            model_name=model_name,
            platform='unknown',  # 需要手动设置
            input_price_per_1k=input_price,
            output_price_per_1k=output_price
        )
        db.session.add(pricing)
    
    try:
        db.session.commit()
        return True
    except Exception as e:
        db.session.rollback()
        print(f"更新模型定价失败: {e}")
        return False

def cleanup_old_records(days=30):
    """清理旧记录（可选的维护功能）"""
    from datetime import timedelta
    from models import RequestRecord
    
    cutoff_date = datetime.utcnow() - timedelta(days=days)
    
    # 删除旧的请求记录（保留最近30天）
    old_requests = db.session.query(RequestRecord).filter(
        RequestRecord.created_at < cutoff_date,
        RequestRecord.status.in_(['completed', 'failed'])
    ).all()
    
    for request in old_requests:
        db.session.delete(request)
    
    try:
        db.session.commit()
        print(f"清理了 {len(old_requests)} 条旧记录")
    except Exception as e:
        db.session.rollback()
        print(f"清理旧记录失败: {e}")